
# 测试模型 将结果在图片中画出来

import os
from mmdet.apis import init_detector, inference_detector, show_result_pyplot
import mmcv
import argparse


def read_image():
    with open("/media/psdz/CAIYQDisk4/code/package/mmdetection/data/coco/val.txt", "r") as fp:
        images = fp.read().split("\n")
        return images


def get_res(config_file, checkpoint_file, model_name):
    prix = '/media/psdz/CAIYQDisk4/code/package/mmdetection/data/coco/train2017_end/'
    model = init_detector(config_file, checkpoint_file, device='cuda:0')
    images = read_image()
    for img_name in images:
        # image路径
        img = prix + img_name + ".jpg"
        result_data = inference_detector(model, img)
        print(img_name)
        # result_img = show_result_pyplot(model, img, result_data)
        if hasattr(model, 'module'):
            model = model.module
        model.show_result(
                        img,
                        result_data,
                        score_thr=0.3,
                        show=False,
                        wait_time=0,
                        win_name='result',
                        bbox_color=(72, 101, 241),
                        text_color=(72, 101, 241),
                        out_file='/media/psdz/CAIYQDisk4/code/package/mmdetection/result/%s/res_' % model_name + img_name + '.jpg')


def parse_args():
    parser = argparse.ArgumentParser(description='Train a detector')
    parser.add_argument('config', help='test config file path')
    parser.add_argument('pth', help='test model pth')
    parser.add_argument('model', help='model name')

    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    config_file = args.config
    checkpoint_file = args.pth
    model_name = args.model
    get_res(config_file, checkpoint_file, model_name=model_name)


def local_main():
    config_file = "../work_dirs/atss_crpn_2_fl_grid_rcnn_r50_fpn_gn-head_1x_coco/" \
                  "atss_crpn_end_grid_rcnn_r50_fpn_gn-head_1x_coco.py"
    checkpoint_file = "../work_dirs/atss_crpn_2_fl_grid_rcnn_r50_fpn_gn-head_1x_coco/" \
                      "epoch_max.pth"
    model_name = "atss_crpn_2fl"
    get_res(config_file, checkpoint_file, model_name)


if __name__ == '__main__':
    main()
    # local_main()